Source code for hysop.tools.mpi_utils

# Copyright (c) HySoP 2011-2024
#
# This file is part of HySoP software.
# See "https://particle_methods.gricad-pages.univ-grenoble-alpes.fr/hysop-doc/"
# for further info.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import numpy as np

from hysop.core.mpi import MPI
from hysop.tools.numerics import get_dtype


[docs] def create_sized(nbytes): mpi_type = MPI.Datatype.Create_contiguous(MPI.BYTE, nbytes) mpi_type.Commit() return mpi_type
[docs] def iter_mpi_requests(requests): """Iterate on MPI events with Waitany.""" while True: idx = MPI.Request.Waitany(requests) if idx == MPI.UNDEFINED: break yield idx
MPI_HALF = create_sized(2)
[docs] def dtype_to_mpi_type(dtype): """Try to convert a numpy dtype to a MPI type.""" if dtype is None: return None from mpi4py import MPI dtype = get_dtype(dtype) mpi_types = { np.int8: MPI.SIGNED_CHAR, np.int16: MPI.SIGNED_SHORT, np.int32: MPI.INT, np.int64: MPI.LONG_LONG, np.uint8: MPI.UNSIGNED_CHAR, np.uint16: MPI.UNSIGNED_SHORT, np.uint32: MPI.UNSIGNED, np.uint64: MPI.UNSIGNED_LONG_LONG, np.float16: MPI_HALF, np.float32: MPI.FLOAT, np.float64: MPI.DOUBLE, np.float128: MPI.LONG_DOUBLE, # (80 bits floats on a 64 bit machine) np.complex64: MPI.COMPLEX, np.complex128: MPI.DOUBLE_COMPLEX, } if hasattr(MPI, "LONG_DOUBLE_COMPLEX"): mpi_types[np.complex256] = MPI.LONG_DOUBLE_COMPLEX if dtype not in mpi_types: msg = f"Unknown dtype {dtype}." raise NotImplementedError(msg) return mpi_types[dtype]
[docs] def order_to_mpi_order(order): from hysop.constants import MemoryOrdering if ( (order in "cC") or (order == MemoryOrdering.C_CONTIGUOUS) or (order == MPI.ORDER_C) ): return MPI.ORDER_C elif ( (order in "fF") or (order == MemoryOrdering.F_CONTIGUOUS) or (order == MPI.ORDER_F) ): return MPI.ORDER_F else: msg = f"Unknown value of type {type(order)}." raise ValueError(msg)
[docs] def get_mpi_order(data): from hysop.core.arrays.array import Array if isinstance(data, Array): is_c_contiguous = data.is_c_contiguous is_f_contiguous = data.is_fortran_contiguous else: # assume numpy like interface is_c_contiguous = data.flags.c_contiguous is_f_contiguous = data.flags.f_contiguous if is_c_contiguous: return MPI.ORDER_C elif is_f_contiguous: return MPI.ORDER_F else: msg = "Data is neither C, nor Fortran contiguous." raise ValueError(msg)